from torch import nn
import torch.nn.functional as F
import time
import torch

class Residual_Block (nn.Module):
    def __init__(self,i_channel,o_channel,stride=1,downsample=None):
        super(Residual_Block,self).__init__()
        self.conv1=nn.Conv2d(in_channels=i_channel,out_channels=o_channel,kernel_size=3,stride=stride,padding=1,bias=False)
        self.bn1=nn.BatchNorm2d(o_channel)
        self.relu=nn.ReLU(inplace=True)
        
        self.conv2=nn.Conv2d(in_channels=o_channel,out_channels=o_channel,kernel_size=3,stride=1,padding=1,bias=False)
        self.bn2=nn.BatchNorm2d(o_channel)
        self.downsample=downsample
        
    def forward(self,x):
        residual=x
        out=self.conv1(x)
        out=self.bn1(out)
        out=self.relu(out)
        out=self.conv2(out)
        out=self.bn2(out)
        if self.downsample:
            residual=self.downsample(x)
        out+=residual
        out=self.relu(out)
        return out


class ResNet(nn.Module):
    def __init__(self,block,layers,num_classes=10):
        super(ResNet,self).__init__()
        self.layer_in_channels=32
        self.conv=nn.Conv2d(in_channels=3,out_channels=32,kernel_size=3,stride=1,padding=1,bias=False)
        # self.maxpool = nn.MaxPool2d(kernel_size=2,stride=2) 
        self.bn=nn.BatchNorm2d(32)
        self.relu=nn.ReLU(inplace=True)
        self.layer1=self.make_layer(block,32,layers[0])#32
        self.layer2=self.make_layer(block,64,layers[1],2) #16
        self.layer3=self.make_layer(block,128,layers[2],2) #8
        self.layer4=self.make_layer(block,256,layers[3],2) #4
        self.max_pool=nn.MaxPool2d(kernel_size=4)
        self.fc=nn.Linear(256,num_classes)
    
    def make_layer(self,block,out_channels,blocks,stride=1):
        downsample=None
        if (stride!=1) or (self.layer_in_channels !=out_channels):
            downsample=nn.Sequential(
                                     nn.Conv2d(self.layer_in_channels,out_channels,kernel_size=3,stride=stride,padding=1,bias=False),
                                     nn.BatchNorm2d(out_channels)
                                     )
        layers=[]
        layers.append(block(self.layer_in_channels,out_channels,stride,downsample))
        self.layer_in_channels=out_channels
        for i in range(1,blocks):
            layers.append(block(out_channels,out_channels))
        return nn.Sequential(*layers)
            
    
    def forward(self,x):
        out = self.conv(x)
        out = self.bn(out)    
        out = self.relu(out)
        out = self.layer1(out) 
        out = self.layer2(out) 
        out = self.layer3(out) 
        out = self.layer4(out) 
        out = self.max_pool(out) 
        out = out.view(x.size()[0], -1) 
        out = self.fc(out)
        return out 
    
    def train(self,train_data,config,optim,device,scheduler=None,):
        
        train_loader = torch.utils.data.DataLoader(train_data,batch_size=config["batchsize"],shuffle = True, pin_memory = True,num_workers = 2)
        Loss_list = []
        Accuracy_list = []
        false = 0
        self.to(device)
        for epo in range(config["epoch"]):
            begin = time.time()
            num_batch = 0
            for img, lab in train_loader:
                num_batch += 1

                img = img.to(device)
                lab = lab.to(device)
                pre_label = self(img)
                loss = F.cross_entropy(input=pre_label, target=lab.long()).mean()
                pre_label = torch.argmax(pre_label, dim=1)
                false = false+(pre_label!=lab).sum()
                self.zero_grad()
                loss.backward()
                optim.step()

                if num_batch >= (len(train_loader)//10):
                    num_batch = 0
                    print("|###|",end="")
            print(" ")
            false = false.detach().cpu().numpy()
            true = (len(train_data)-false)/len(train_data)
            print("epoch:{}  Accuracy:{:.2f}%  Loss:{}".format(epo,true*100,loss.detach().cpu().numpy()),end="")
            false = 0
            path = "parameters/regularization/model_parameter_lr{}_epoch{}.pt".format(config["lr"],epo)
            torch.save(self.state_dict(),path)
            Loss_list.append(loss.detach().cpu().numpy())
            Accuracy_list.append(true)
            end = time.time()
            print(f"time:{end-begin}")
            
        return Accuracy_list, Loss_list


    def run(self,test_data,device):
        torch.no_grad()
        test_loader = torch.utils.data.DataLoader(test_data,batch_size=32,shuffle = True, pin_memory = True,num_workers = 2)
        false = 0
        num_batch = 0
        self.to(device)
        for img, lab in test_loader:
            num_batch = num_batch+1 
            img = img.to(device)
            lab = lab.to(device)
            pre_label = self(img)
            pre_label = torch.argmax(pre_label, dim=1)

            false = false+(pre_label!=lab).sum()
            if num_batch >= (len(test_loader)//10):
                num_batch = 0
                print("|###|",end="")
        print("")
        false = false.detach().cpu().numpy()
        true = (len(test_data)-false)/len(test_data)
        print("正确率:{:.2f}%".format(true*100))
        
        return true

